
use std::io::{ Read, BufReader };
use std::fs::File;
use std::collections::HashMap;

use crate::jpeg_meta::*;
use crate::jpeg_utils::*;

trait JpegMarkerSegmentReader {
    fn read_app0(&mut self) -> App0Info;

    fn read_dqt(&mut self) -> (u8, [f32;64]);

    fn read_sof0(&mut self) -> SOFInfo;

    fn read_dht(&mut self) -> (u8, u8, HashMap<(u8,u16),u8>);

    fn read_sos(&mut self) -> SOSInfo;
}

struct JpegBitReader<'a> {
    pub count: u8,          // The number of bits remaining in a byte read
    pub data: u8,           // The byte data
    pub dc: [f32;3],        // The dc coffience data
    pub reader: &'a mut BufReader<File>,
}

enum AcCode {
    AllZeros,
    SixteenZeros,
    Normal {
        zeros: u8,
        ac: f32,
    }
}

impl<'a> JpegBitReader<'a> {
    pub fn new(reader: &'a mut BufReader<File>) -> Self {
        JpegBitReader {
            count: 0,
            data: 0,
            dc: [0f32; 3],
            reader,
        }
    }

    pub fn read_mcus(&mut self, jpeg_meta: &JPEGMetaData) -> Vec<Vec<MCUInfo>> {
        let sof = &jpeg_meta.sof;

        let mcu_height = 8 * sof.v_max_samples as u16;
        let mcu_width = 8 * sof.h_max_samples as u16;

        let mcu_row = ((sof.height - 1) / mcu_height + 1) as usize;
        let mcu_col = ((sof.width - 1) / mcu_width + 1) as usize;

        let mut mcus:Vec<Vec<MCUInfo>> = vec![vec![MCUInfo::default(); mcu_col]; mcu_row];
        for h in 0..mcu_row {
            for w in 0..mcu_col {
                mcus[h][w] = self.read_mcu(&jpeg_meta);
            }
        }
        mcus
    }
}

impl<'a> JpegBitReader<'a> {

    fn read_bit(&mut self) -> u8 {
        if self.count == 0 {
            self.data = jpeg_utils::read_byte(self.reader);
            if self.data == 0xFF {
                let check = jpeg_utils::read_byte(self.reader);
                assert_eq!( check, 0x0 );
            }
        }
        let val = if self.data & (0x1 << (7 - self.count)) > 0 { 1 } else { 0 };
        self.count = (self.count + 1) % 8;
        val
    }

    fn read_value(&mut self, n: u8) -> f32 {
        assert!(n >= 1);

        let first = self.read_bit();
        let mut code = 1;
        for _ in 1..n {
            let val = self.read_bit();
            code <<= 1;
            code += if val == first { 1 } else { 0 };
        }
        if first == 0 {
            -code as f32
        } else {
            code as f32
        }
    }

    fn read_mcu(&mut self, jpeg_meta: &JPEGMetaData) ->MCUInfo {
        let sof = &jpeg_meta.sof;
        let mut mcu = MCUInfo::default();
        for id in 0..sof.num_com {
            mcu[id as usize] = self.read_color_component_block(id, &jpeg_meta);
        }
        mcu
    }

    fn read_color_component_block(&mut self, id: u8, jpeg_meta: &JPEGMetaData) -> ComBlockInfo {
        let sof = &jpeg_meta.sof;
        let row = sof.coms[id as usize].v_sample_factor as usize;
        let col = sof.coms[id as usize].h_sample_factor as usize;

        let mut blocks = vec![vec![BlockInfo::default(); col]; row];
        for r in 0..row {
            for c in 0..col {
                blocks[r][c] = self.read_block(id, &jpeg_meta);
            }
        }

        blocks
    }

    fn read_block(&mut self, id: u8, jpeg_meta: &JPEGMetaData) -> BlockInfo {
        let sos = &jpeg_meta.sos.coms[id as usize];
        let dc_table = &jpeg_meta.huffmantable.dc_table[sos.dc_id as usize];
        let ad_table = &jpeg_meta.huffmantable.ac_table[sos.ac_id as usize];
        let mut block = BlockInfo::default();
        block[0][0] = self.read_dc(&dc_table, id);
        let mut idx = 1;
        while idx < 64 {
            let ac = self.read_ac(&ad_table);
            match ac {
                AcCode::AllZeros => {
                    while idx < 64 {
                        block[idx / 8][idx % 8] = 0.0f32;
                        idx += 1;
                    }
                },
                AcCode::SixteenZeros => {
                    for _ in 0..16 {
                        block[idx / 8][idx % 8] = 0.0f32;
                        idx += 1;
                    }
                }
                AcCode::Normal {zeros, ac} => {
                    for _ in 0..zeros {
                        block[idx / 8][idx % 8] = 0.0f32;
                        idx += 1;
                    }
                    block[idx / 8][idx % 8] = ac;
                    idx += 1;
                }
            }
        }
        block
    }

    fn match_huffmantable(&mut self, table: &HashMap<(u8,u16), u8>) -> u8 {
        let mut code = 0u16;
        let mut len = 1;
        loop {
            if len > 16 {
                panic!("No matching huffman table was found!");
            }

            code += self.read_bit() as u16;
            if table.contains_key(&(len, code)) {
                return *table.get(&(len, code)).unwrap();
            }

            len += 1;
            code = code << 1;
        }
    }

    fn read_dc(&mut self, dc_table: &HashMap<(u8,u16),u8>, id: u8) -> f32 {
        let code_len = self.match_huffmantable(dc_table);
        if code_len == 0 {
            return self.dc[id as usize];
        }
        self.dc[id as usize] += self.read_value(code_len);
        return self.dc[id as usize];
    }

    fn read_ac(&mut self, ac_table: &HashMap<(u8,u16),u8>) -> AcCode {
        let code_len = self.match_huffmantable(ac_table);
        match code_len {
            0x00 => AcCode::AllZeros,
            0xF0 => AcCode::SixteenZeros,
            x => {
                AcCode::Normal {
                    zeros: x >> 4,
                    ac: self.read_value(x & 0x0F),
                }
            }
        }
    }
}

impl JpegMarkerSegmentReader for BufReader<File> {

    fn read_app0(mut self: &mut Self) -> App0Info {
        let mut size = jpeg_utils::read_word(&mut self);

        let mut app = App0Info::default();

        self.read_exact(&mut app.identifier).expect("Failed to read the identifier of the jpeg");
        self.read_exact(&mut app.version).expect("Failed to read the version of the jpeg");
        app.density_units = jpeg_utils::read_byte(&mut self);
        app.x_density = jpeg_utils::read_word(&mut self);
        app.y_density = jpeg_utils::read_word(&mut self);
        app.x_thumbnail = jpeg_utils::read_byte(&mut self);
        app.y_thumbnail = jpeg_utils::read_byte(&mut self);

        size -= 16;
        assert_eq!(size, 0);

        app
    }

    fn read_dqt(mut self: &mut Self) -> (u8, [f32;64]) {
        let mut size = jpeg_utils::read_word(&mut self);

        let precision_quantable = jpeg_utils::read_byte(&mut self);
        let precision = precision_quantable >> 4;
        let quantable = precision_quantable & 0x0F;

        let mut table: [f32; 64] = [0f32; 64];
        if precision == 0 {
            for i in 0..64 {
                table[i] = jpeg_utils::read_byte(&mut self) as f32;
            }

            size -= 67;
        } else if precision == 1 {
            for i in 0..64 {
                table[i] = jpeg_utils::read_word(&mut self) as f32;
            }

            size -= 131;
        } else {
            panic!("The Invalid Quantable Precision!");
        }

        assert_eq!(size, 0);
        (quantable, table)
    }

    fn read_sof0(mut self: &mut Self) -> SOFInfo {
        let mut size = jpeg_utils::read_word(&mut self);

        let mut sof = SOFInfo::default();

        sof.precision_com = jpeg_utils::read_byte(&mut self);
        sof.height = jpeg_utils::read_word(&mut self);
        sof.width = jpeg_utils::read_word(&mut self);
        sof.num_com = jpeg_utils::read_byte(&mut self);
        for _ in 0..sof.num_com {
            let mut com = ComInfoForSOF::default();
            com.com_id = jpeg_utils::read_byte(&mut self);
            let sample = jpeg_utils::read_byte(&mut self);
            com.h_sample_factor = sample >> 4;
            com.v_sample_factor = sample & 0x0F;
            if com.h_sample_factor > sof.h_max_samples {
                sof.h_max_samples = com.h_sample_factor;
            }
            if com.v_sample_factor > sof.v_max_samples {
                sof.v_max_samples = com.v_sample_factor;
            }
            com.quantable_id = jpeg_utils::read_byte(&mut self);

            sof.coms.push(com);
        }

        size -= 8 + 3 * sof.num_com as u16;
        assert_eq!(size, 0);

        sof
    }

    fn read_dht(mut self: &mut Self) -> (u8, u8, HashMap<(u8,u16),u8>) {
        let mut size = jpeg_utils::read_word(&mut self);

        let ttype_id = jpeg_utils::read_byte(&mut self);
        let table_type = ttype_id >> 4;
        let table_id = ttype_id & 0x0F;

        let mut code_num = [0u8; 16];
        self.read_exact(&mut code_num).expect("Faild to read the number of Huffman codes of each length");

        let mut table: HashMap<(u8,u16),u8> = Default::default();

        size -= 19;

        let mut code = 0;
        for i in 0..16 {
            for _ in 0..code_num[i] {
                let val = jpeg_utils::read_byte(&mut self);
                table.insert((i as u8 + 1, code), val);
                code += 1;
                size -= 1;
            }
            code <<= 1;
        }

        assert_eq!(size, 0);

        (table_type, table_id, table)
    }

    fn read_sos(mut self: &mut Self) -> SOSInfo {
        let mut size = jpeg_utils::read_word(&mut self);

        let mut sos = SOSInfo::default();
        sos.num_com = jpeg_utils::read_byte(&mut self);
        for _ in 0..sos.num_com {
            let mut com = ComInfoForSOS::default();
            com.com_id = jpeg_utils::read_byte(&mut self);
            let dc_ac = jpeg_utils::read_byte(&mut self);
            com.dc_id = dc_ac >> 4;
            com.ac_id = dc_ac & 0x0F;
            sos.coms.push(com);
        }
        jpeg_utils::read_byte(&mut self);
        jpeg_utils::read_byte(&mut self);
        jpeg_utils::read_byte(&mut self);

        size -= 6 + 2 * sos.num_com as u16;
        assert_eq!(size, 0);

        sos
    }

}

pub mod jpeg_read {
    use super::*;

    pub fn data_read(mut reader: &mut BufReader<File>) -> (JPEGMetaData, Vec<Vec<MCUInfo>>) {
        let mut jpeg_meta = JPEGMetaData::new();

        let mut mcu_info = Default::default();

        loop {
            let tag = jpeg_utils::read_byte(&mut reader);
            if tag != 0xFF {
                continue;
            }
            let tag = jpeg_utils::read_byte(&mut reader);
            match tag.into() {
                SegmentTag::SOI => continue,
                SegmentTag::EOI => break,
                SegmentTag::APP0 => {
                    let app = reader.read_app0();
                    println!("The Application Segment: {:#?}", app);
                },
                SegmentTag::DQT => {
                    let (id, table) = reader.read_dqt();
                    jpeg_meta.quantable[id as usize] = table;
                },
                SegmentTag::SOF0 => {
                    jpeg_meta.sof = reader.read_sof0();
                },
                SegmentTag::SOS => {
                    jpeg_meta.sos = reader.read_sos();
                    let mut bit_reader = JpegBitReader::new(&mut reader);
                    mcu_info = bit_reader.read_mcus(&jpeg_meta);
                }
                SegmentTag::DHT => {
                    let (ttyp, tid, table) = reader.read_dht();
                    if ttyp == 0x0 {
                        jpeg_meta.huffmantable.dc_table[tid as usize] = table;
                    } else if ttyp == 0x1 {
                        jpeg_meta.huffmantable.ac_table[tid as usize] = table;
                    } else {
                        panic!("The Invalid Huffmant Type: {}", ttyp);
                    }
                }
                _ => {
                    continue
                }
            }
        }

        (jpeg_meta, mcu_info)
    }
}
